Skip to main content

Partition to K Equal Sum Subsets

Question​

Given an array of integers nums and a positive integer k, divide the array into k non-empty subsets whose sums are all equal.

Example 1
Input: nums = [4, 3, 2, 3, 5, 2, 1], k = 4

Output: True 
Subsets: [4, 3, 2], [3, 5], [2, 1], [6]

Solution​

all//Partition to K Equal Sum Subsets.py


# given an array of integers and a number K, the task is to divide the array into K non-empty subsets whose sum is equal
#partition to k equal sum subsets

def partition_to_k_subsets(nums, k):
    # check if the total sum can be divided by k
    total_sum = sum(nums)
    if total_sum % k != 0:
        return False

    # calculate the required sum for each subset
    subset_sum = total_sum // k
    # sort the array
    nums.sort()

    # array to hold the subsets
    subsets = [[] for _ in range(k)]

    # index used to fill the subsets
    index = 0

    # traverse the array from the end
    for i in range(len(nums)-1, -1, -1):
        # if the current element can fit the current subset
        if nums[i] <= subset_sum - sum(subsets[index]):
            # add the current element
            subsets[index].append(nums[i])
        else:
            # move to the next subset
            index += 1
            # add the current element
            subsets[index].append(nums[i])

    # check if all the subsets are filled
    if index == k-1:
        return subsets
    else:
        return False